# -*- coding: utf-8 -*-
# @Time    : 2021/3/11 18:14
# @Author  : huangwei
# @File    : get_signature.py
# @Software: PyCharm
import os
from get_signature_method import *
from skimage import measure

# 参数设置
SIZE = 512, 512  # 模型输入尺寸
WEIGHT_PATH = "models/table.weights"

tableNet = cv2.dnn.readNetFromDarknet(WEIGHT_PATH.replace('.weights', '.cfg'), WEIGHT_PATH)


def table_predict( img ):
    """ 将图片传入模型中得到输出 """
    imgResize, fx, fy, dx, dy = resize_img(img, SIZE)

    imgResize = np.array(imgResize)
    img_w, img_h = SIZE

    # 对图像进行减均值、缩放以及交换通道
    image = cv2.dnn.blobFromImage(imgResize, 1, size=(img_w, img_h), swapRB=False)

    image = np.array(image) / 255
    tableNet.setInput(image)
    out = tableNet.forward()
    out = exp(out[0])
    out = out[:, dy:, dx:]
    return out, fx, fy


def get_table_rowcols( img, prob, row=100, col=100 ):
    out, fx, fy = table_predict(img)

    rows = out[0]
    cols = out[1]

    # 实现连通区域标记，connectivity=2 表示8连通，即该点周围8个点都与其值相等，该处标记则为True
    labels = measure.label(rows > prob, connectivity=2)
    # 返回所有连通区域的属性列表
    regions = measure.regionprops(labels)
    # 该连通区域的边界外接框宽度大于阈值row则将该外接框的横向中心线加入到集合中
    RowsLines = [get_line(line.coords) for line in regions if line.bbox[3] - line.bbox[1] > row]

    # 同理找出所有竖线
    labels = measure.label(cols > prob, connectivity=2)
    regions = measure.regionprops(labels)
    ColsLines = [get_line(line.coords) for line in regions if line.bbox[2] - line.bbox[0] > col]

    # 将所有找出的线在黑色图片上画出，找出只有这些线存在的图片上的连通区域
    tmp = np.zeros(SIZE[::-1], dtype='uint8')
    tmp = draw_lines(tmp, ColsLines + RowsLines, color=255, line_width=1)
    labels = measure.label(tmp > 0, connectivity=2)
    regions = measure.regionprops(labels)

    # 去除这些连通区域中有一边较小的区域以及不在整个大连通区域的线段
    for region in regions:
        ymin, xmin, ymax, xmax = region.bbox
        label = region.label
        if ymax - ymin < 20 or xmax - xmin < 20:
            labels[labels == label] = 0
    labels = measure.label(labels > 0, connectivity=2)

    indY, indX = np.where(labels > 0)
    xmin, xmax = indX.min(), indX.max()
    ymin, ymax = indY.min(), indY.max()
    RowsLines = [p for p in RowsLines if
                 xmin <= p[0] <= xmax and xmin <= p[2] <= xmax and ymin <= p[1] <= ymax and ymin <= p[3] <= ymax]
    ColsLines = [p for p in ColsLines if
                 xmin <= p[0] <= xmax and xmin <= p[2] <= xmax and ymin <= p[1] <= ymax and ymin <= p[3] <= ymax]

    # 将坐标转换成原图中的坐标
    RowsLines = [[box[0] / fx, box[1] / fy, box[2] / fx, box[3] / fy] for box in RowsLines]
    ColsLines = [[box[0] / fx, box[1] / fy, box[2] / fx, box[3] / fy] for box in ColsLines]
    return RowsLines, ColsLines


def adjust_lines( RowsLines, ColsLines, alph=50 ):
    """找出某些舍去了的短的线段"""
    nrow = len(RowsLines)
    ncol = len(ColsLines)
    newRowsLines = []
    newColsLines = []

    for i in range(nrow):
        x1, y1, x2, y2 = RowsLines[i]
        cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
        for j in range(nrow):
            if i != j:
                x3, y3, x4, y4 = RowsLines[j]
                cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
                if (x3 < cx1 < x4 or y3 < cy1 < y4) or (x1 < cx2 < x2 or y1 < cy2 < y2):
                    continue
                else:
                    r = dist((x1, y1), (x3, y3))
                    if r < alph:
                        newRowsLines.append([x1, y1, x3, y3])
                    r = dist((x1, y1), (x4, y4))
                    if r < alph:
                        newRowsLines.append([x1, y1, x4, y4])

                    r = dist((x2, y2), (x3, y3))
                    if r < alph:
                        newRowsLines.append([x2, y2, x3, y3])
                    r = dist((x2, y2), (x4, y4))
                    if r < alph:
                        newRowsLines.append([x2, y2, x4, y4])

    for i in range(ncol):
        x1, y1, x2, y2 = ColsLines[i]
        cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
        for j in range(ncol):
            if i != j:
                x3, y3, x4, y4 = ColsLines[j]
                cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
                if (x3 < cx1 < x4 or y3 < cy1 < y4) or (x1 < cx2 < x2 or y1 < cy2 < y2):
                    continue
                else:
                    r = dist((x1, y1), (x3, y3))
                    if r < alph:
                        newColsLines.append([x1, y1, x3, y3])
                    r = dist((x1, y1), (x4, y4))
                    if r < alph:
                        newColsLines.append([x1, y1, x4, y4])

                    r = dist((x2, y2), (x3, y3))
                    if r < alph:
                        newColsLines.append([x2, y2, x3, y3])
                    r = dist((x2, y2), (x4, y4))
                    if r < alph:
                        newColsLines.append([x2, y2, x4, y4])

    return newRowsLines, newColsLines


def get_cell( img, prob, alph, row=100, col=100 ):
    """
    从图片中获取所有的单元格
    """
    w, h = SIZE
    RowsLines, ColsLines = get_table_rowcols(img, prob, row, col)
    newRowsLines, newColsLines = adjust_lines(RowsLines, ColsLines, alph=alph)
    RowsLines += newRowsLines
    ColsLines += newColsLines

    nrow = len(RowsLines)
    ncol = len(ColsLines)

    for i in range(nrow):
        for j in range(ncol):
            RowsLines[i] = line_line(RowsLines[i], ColsLines[j], 32)
            ColsLines[j] = line_line(ColsLines[j], RowsLines[i], 32)

    # 将找到的所有线都在全黑的图上画出，再在这个图上寻找连通区域即每一个格子
    tmp = np.zeros((img.size[1], img.size[0]), dtype='uint8')
    tmp = draw_lines(tmp, ColsLines + RowsLines, color=255, line_width=1)
    tabelLabels = measure.label(tmp == 0, connectivity=2)
    regions = measure.regionprops(tabelLabels)

    rboxes = []
    for region in regions:
        if region.bbox_area < h * w - 10:
            rbox = minAreaRectBox(region.coords)
            rboxes.append(rbox)

    return rboxes


def get_signature( image_path ):
    img = Image.open(image_path).convert('RGB')
    # 获取单元格
    boxes = get_cell(img, prob=0.5, alph=10)

    # 根据 boxes 的 Y 轴坐标将其划分
    center_y = []
    for box in boxes:
        x1, y1, x2, y2, x3, y3, x4, y4 = box
        cy = (y1 + y2 + y3 + y4) / 4
        center_y.append(cy)

    # 按照中心点的 y 坐标的大小进行重新排序
    center_y = np.array(center_y)
    sort_index = np.argsort(center_y)

    boxes = np.array(boxes)
    boxes = boxes[sort_index]
    boxes = boxes.tolist()
    center_y = center_y[sort_index]
    center_y = center_y.tolist()

    # 判断从该点开始接下来四个点是否都在一条直线上
    i = 0
    j = 0
    while i < len(center_y) - 3:
        tmp_y = center_y[i]
        i += 1
        if center_y[i] - tmp_y < 10:
            i += 1
            if center_y[i] - tmp_y < 10:
                i += 1
                if center_y[i] - tmp_y < 10:
                    # 找到四个box了，需要对四个框从左到右排序
                    get_name(img, boxes[i - 3:i + 1], j)
                    j += 1

    print("this picture split signatures finish!")


if __name__ == '__main__':
    img_path = "image/test.jpg"
    if os.path.exists(img_path):
        get_signature(img_path)
    else:
        print("please check your filepath, the file is not exists")
